from pyspark.sql import SparkSession,functions as F
from pyspark.sql.types import *
import pandas as pd

# 代码中需要连接hive读取hive中的表数据
# spark.sql.warehouse.dir 指定数仓位置
# hive.metastore.uris 指定连接hive的metastore服务
spark = SparkSession.builder.\
    config('spark.sql.warehouse.dir','hdfs://node1:8020/user/hive/warehouse').\
    config('hive.metastore.uris','thrift://node1:9083').\
    config('spark.sql.execution.arrow.pyspark.enabled','true').\
    enableHiveSupport().\
    getOrCreate()

@F.pandas_udf(returnType=IntegerType())
def func(sex:pd.Series,dx:pd.Series,lx:pd.Series) -> int:
    # sex整列数据  接受dx整列数据  ，lx的整列数据
    print('sex的值')
    print(sex)
    print('dx的值')
    print(dx)
    print('lx的值')
    print(lx)
    print('*'*100)
    return 1

# 使用pandas的Series 可以接受整列字段数据
@F.pandas_udf(returnType=FloatType())
def func2(dx:pd.Series, lx:pd.Series) -> float:
    # 接收 dx整列数据，lx整列数据
    dx_tmp = 0
    lx_tmp = 0
    for i in range(0,len(lx)): # 通过len获取lx元素遍历次数
        if i== 0:  # 条件成立，说明是第一行数据
            dx_tmp = dx[i]   # 获取第一行的 dx数据
            lx_tmp = lx[i]   # 获取第一行的 lx数据
        else:
            # 说明是后几行数据    开始计算第二行往后的数据
            lx_tmp = (dx_tmp + lx_tmp)/2   # 计算的是后几行数据
            dx_tmp = dx[i]   # 获取dx_tmp的后几行数据
    print(lx_tmp)
    return lx_tmp


# 注册自定义函数
spark.udf.register('func',func)

# 进行sql使用
df =spark.sql('select sex, policy_year,dx,lx, func2(dx, lx) over(partition by sex order by policy_year) as lx_new from tb2_1')
df.show()